import numpy as np
import stl
import math
import vtkplotlib as vpl
import json
import time


def rotate_x(vector, angle):
    angle = math.radians(angle)
    rotate_matrix = np.array(
        [[1, 0, 0], [0, math.cos(angle), -math.sin(angle)], [0, math.sin(angle), math.cos(angle)]])
    rotate_vector = np.dot(vector, rotate_matrix)
    return rotate_vector


def rotate_y(vector, angle):
    angle = math.radians(angle)
    rotate_matrix = np.array(
        [[math.cos(angle), 0, math.sin(angle)], [0, 1, 0], [-math.sin(angle), 0, math.cos(angle)]])
    rotate_vector = np.dot(vector, rotate_matrix)
    return rotate_vector


def rotate_z(vector, angle):
    angle = math.radians(angle)
    rotate_matrix = np.array(
        [[math.cos(angle), -math.sin(angle), 0], [math.sin(angle), math.cos(angle), 0], [0, 0, 1]])
    rotate_vector = np.dot(vector, rotate_matrix)
    return rotate_vector


class StlMesh:
    def __init__(self, path):
        self.mesh = stl.Mesh.from_file(path)
        self.x = self.mesh.x
        self.y = self.mesh.y
        self.z = self.mesh.z
        self.vectors = self.mesh.vectors
        # self.triangles_data = np.array([np.array(vector) for vector in self.vectors])
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.id = None
        self.name = None
        self.data = self.mesh.data
        self.pos = np.array([0,0,0], dtype=np.float32)
        self.vel = np.array([0,0,0], dtype=np.float32)
        self.acc = np.array([0,0,0], dtype=np.float32)

    def move(self, dx, dy, dz):
        self.x += dx
        self.y += dy
        self.z += dz
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])

    def scale(self, kx, ky, kz):
        self.x *= kx
        self.y *= ky
        self.z *= kz
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])

    def rotate(self, ax, ay, az):
        # start_time = time.time()
        cog0 = self.cog
        self.move_to_centre()
        self.vectors[:, :] = rotate_x(self.vectors[:, :], ax)
        self.vectors[:, :] = rotate_y(self.vectors[:, :], ay)
        self.vectors[:, :] = rotate_z(self.vectors[:, :], az)
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.move(cog0[0], cog0[1], cog0[2])
        # print('Time rotate', time.time()-start_time)

    def move_to_centre(self):
        # volume, cog, inertia = self.mesh.get_mass_properties()
        self.move(-self.cog[0], -self.cog[1], -self.cog[2])

    def show_mesh(self, fig, axes=True):
        vpl.mesh_plot(self.mesh, fig=fig)
        if axes:
            x_ax = np.array([[0, 0, 0], [np.max(self.x) * 1.2, 0, 0]])
            y_ax = np.array([[0, 0, 0], [0, np.max(self.y) * 1.2, 0]])
            z_ax = np.array([[0, 0, 0], [0, 0, np.max(self.z) * 1.2]])
            vpl.plot(x_ax, color='red', fig=fig, line_width=5.0)
            vpl.plot(y_ax, color='green', fig=fig, line_width=5.0)
            vpl.plot(z_ax, color='blue', fig=fig, line_width=5.0)


class JsonMesh(StlMesh):
    def __init__(self, scene_json, key):
        self.vectors = np.asarray(scene_json[key]['polygons'], dtype=np.float32)
        self.pos = np.asarray(scene_json[key]['position'], dtype=np.float32)
        self.vel = np.asarray(scene_json[key]['velocity'], dtype=np.float32)
        self.acc = np.asarray(scene_json[key]['acceleration'], dtype=np.float32)
        self.name = key
        self.x = self.vectors[:, :, 0]
        self.y = self.vectors[:, :, 1]
        self.z = self.vectors[:, :, 2]
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2], dtype=np.float32)

    def show_mesh(self, fig, axes=True):
        vpl.mesh_plot(self.vectors, fig=fig)
        if axes:
            x_ax = np.array([[0, 0, 0], [np.max(self.x) * 1.2, 0, 0]])
            y_ax = np.array([[0, 0, 0], [0, np.max(self.y) * 1.2, 0]])
            z_ax = np.array([[0, 0, 0], [0, 0, np.max(self.z) * 1.2]])
            vpl.plot(x_ax, color='red', fig=fig, line_width=5.0)
            vpl.plot(y_ax, color='green', fig=fig, line_width=5.0)
            vpl.plot(z_ax, color='blue', fig=fig, line_width=5.0)


class CubeMesh(StlMesh):
    def __init__(self, left_bottom, right_top):
        self.vertex = np.zeros((8, 3), dtype=np.float32)
        self.vertex[0] = left_bottom
        self.vertex[1] = [left_bottom[0], right_top[1], left_bottom[2]]
        self.vertex[2] = [right_top[0], right_top[1], left_bottom[2]]
        self.vertex[3] = [right_top[0], left_bottom[1], left_bottom[2]]
        self.vertex[4] = [left_bottom[0], left_bottom[1], right_top[2]]
        self.vertex[5] = [left_bottom[0], right_top[1], right_top[2]]
        self.vertex[6] = right_top
        self.vertex[7] = [right_top[0], left_bottom[1], right_top[2]]
        self.vectors = np.array([[self.vertex[0], self.vertex[4], self.vertex[1]],
                              [self.vertex[1], self.vertex[4], self.vertex[5]],
                              [self.vertex[4], self.vertex[7], self.vertex[5]],
                              [self.vertex[5], self.vertex[7], self.vertex[6]],
                              [self.vertex[3], self.vertex[7], self.vertex[6]],
                              [self.vertex[2], self.vertex[3], self.vertex[6]],
                              [self.vertex[0], self.vertex[2], self.vertex[3]],
                              [self.vertex[0], self.vertex[1], self.vertex[2]],
                              [self.vertex[1], self.vertex[6], self.vertex[2]],
                              [self.vertex[1], self.vertex[5], self.vertex[6]],
                              [self.vertex[0], self.vertex[7], self.vertex[4]],
                              [self.vertex[0], self.vertex[3], self.vertex[7]]])
        print(self.vectors)
        self.x = self.vectors[:, :, 0]
        self.y = self.vectors[:, :, 1]
        self.z = self.vectors[:, :, 2]
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.pos = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.vel = np.array([0, 0, 0], dtype=np.float32)
        self.acc = np.array([0, 0, 0], dtype=np.float32)

    def show_mesh(self, fig):
        vpl.mesh_plot(self.vectors, fig=fig)
        x_ax = np.array([[0, 0, 0], [np.max(self.x) * 1.2, 0, 0]])
        y_ax = np.array([[0, 0, 0], [0, np.max(self.y) * 1.2, 0]])
        z_ax = np.array([[0, 0, 0], [0, 0, np.max(self.z) * 1.2]])
        vpl.plot(x_ax, color='red', fig=fig, line_width=5.0)
        vpl.plot(y_ax, color='green', fig=fig, line_width=5.0)
        vpl.plot(z_ax, color='blue', fig=fig, line_width=5.0)


class PlaneMesh(StlMesh):
    def __init__(self, left_bottom, right_top):
        self.vertex = np.zeros((4, 3), dtype=np.float32)
        self.vertex[0] = left_bottom
        self.vertex[1] = [left_bottom[0], right_top[1], left_bottom[2]]
        self.vertex[2] = right_top
        self.vertex[3] = [right_top[0], left_bottom[1], right_top[2]]
        self.vectors = np.array([[self.vertex[0], self.vertex[1], self.vertex[2]],
                              [self.vertex[0], self.vertex[3], self.vertex[2]]])
        self.x = self.vectors[:, :, 0]
        self.y = self.vectors[:, :, 1]
        self.z = self.vectors[:, :, 2]
        self.cog = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.pos = np.array([(np.max(self.x) + np.min(self.x)) / 2, (np.max(self.y) + np.min(self.y)) / 2,
                             (np.max(self.z) + np.min(self.z)) / 2])
        self.vel = np.array([0, 0, 0], dtype=np.float32)
        self.acc = np.array([0, 0, 0], dtype=np.float32)

    def show_mesh(self, fig):
        vpl.mesh_plot(self.vectors, fig=fig)
        # x_ax = np.array([[0, 0, 0], [np.max(self.x) * 1.2, 0, 0]])
        # y_ax = np.array([[0, 0, 0], [0, np.max(self.y) * 1.2, 0]])
        # z_ax = np.array([[0, 0, 0], [0, 0, np.max(self.z) * 1.2]])
        # vpl.plot(x_ax, color='red', fig=fig, line_width=5.0)
        # vpl.plot(y_ax, color='green', fig=fig, line_width=5.0)
        # vpl.plot(z_ax, color='blue', fig=fig, line_width=5.0)



def to_json(mesh, json_dict):
    polygons = mesh.vectors.tolist()
    position = mesh.pos.tolist()
    velocity = mesh.vel.tolist()
    acceleration = mesh.acc.tolist()
    json_dict[mesh.name] = {'polygons': polygons, 'position': position, 'velocity': velocity,
                            'acceleration': acceleration}


def save_json(json_dict, name_file):
    with open(name_file, 'w') as outfile:
        json.dump(json_dict, outfile)
